{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 12,
   "id": "1fc430ca",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Original message: b'WTF, Alice!'\n",
      "Encryption C1: (725374226968432448365528450463540667799630567526402423184017492362194410232185198007920253993995187216521721662343, 2220163884092685170597202150534200353558680069175326450716109816259628864987790756986054551188110826100753535572404, 578944078995539478898067346383698208062908690132418954413632161636084423626880593706141384237810614307828491209731)\n",
      "Encryption C2: b'\\xdd\\xcfGS\\xfa:j\\xc2&@\\x13'\n",
      "Decrypted message: b'WTF, Alice!'\n"
     ]
    }
   ],
   "source": [
    "from py_ecc.optimized_bls12_381 import G1, G2, multiply, add, pairing, curve_order, FQ2, FQ12\n",
    "from py_ecc.bls.hash_to_curve import hash_to_G2\n",
    "from hashlib import sha256\n",
    "import os\n",
    "\n",
    "# 将 message 哈希为 G2 曲线上的点\n",
    "def H1(message, dst=b'WTF-HASH-TO-CURVE-bls12-381'):\n",
    "    # The DST parameter is domain separation tag for the hash function\n",
    "    point_projective = hash_to_G2(message, dst, sha256)\n",
    "    return point_projective\n",
    "\n",
    "# 将 fq12 element (pairing结果) 哈希为 bytes\n",
    "def H2(fq12_elem):\n",
    "    # 将FQ12元素的所有系数转换为字符串并连接\n",
    "    concatenated_coeffs = ''.join(str(coeff) for coeff in fq12_elem.coeffs)\n",
    "    # 对连接后的字符串进行SHA256哈希处理\n",
    "    hash_digest = sha256(concatenated_coeffs.encode()).digest()\n",
    "    return hash_digest\n",
    "\n",
    "# 计算两个 bytes 的 xor\n",
    "def xor_bytes(bytes1, bytes2):\n",
    "    \"\"\"XOR two bytes objects together.\"\"\"\n",
    "    return bytes(a ^ b for a, b in zip(bytes1, bytes2))\n",
    "\n",
    "# IBE 加密\n",
    "def ibe_encrypt(master_pk, user_id, message):\n",
    "    r = int.from_bytes(os.urandom(32), 'big') % curve_order\n",
    "    \n",
    "    # Compute shared secret using pairing\n",
    "    P_user = H1(user_id.encode())\n",
    "    shared_secret = pairing(P_user, master_pk)\n",
    "    C1 = multiply(G1, r)\n",
    "    C2 = xor_bytes(message, H2(shared_secret ** r))\n",
    "\n",
    "    return C1, C2\n",
    "\n",
    "# IBE 解密\n",
    "def ibe_decrypt(user_sk, C1, C2):\n",
    "    shared_secret = pairing(user_sk, C1)\n",
    "    M = xor_bytes(C2, H2(shared_secret))\n",
    "    return M\n",
    "\n",
    "\n",
    "# 主密钥生成\n",
    "master_sk = int.from_bytes(os.urandom(32), 'big') % curve_order\n",
    "master_pk = multiply(G1, master_sk)\n",
    "\n",
    "# 根据 Alice id 生成私钥\n",
    "alice_id = \"alice@wtf.email\"\n",
    "alice_sk = multiply(H1(alice_id.encode()), master_sk)\n",
    "\n",
    "# Bob 编码消息\n",
    "bob_id = \"bob@wtf.email\"\n",
    "message = b\"WTF, Alice!\"\n",
    "print(\"Original message:\", message)\n",
    "C1, C2 = ibe_encrypt(master_pk, alice_id, message)\n",
    "print(\"Encryption C1:\", C1)\n",
    "print(\"Encryption C2:\", C2)\n",
    "\n",
    "# Alice 解码消息\n",
    "decrypted_message = ibe_decrypt(alice_sk, C1, C2)\n",
    "print(\"Decrypted message:\", decrypted_message)\n",
    "\n",
    "# 示例输出\n",
    "# Original message: b'WTF, Alice!'\n",
    "# Encryption C1: (725374226968432448365528450463540667799630567526402423184017492362194410232185198007920253993995187216521721662343, 2220163884092685170597202150534200353558680069175326450716109816259628864987790756986054551188110826100753535572404, 578944078995539478898067346383698208062908690132418954413632161636084423626880593706141384237810614307828491209731)\n",
    "# Encryption C2: b'\\xdd\\xcfGS\\xfa:j\\xc2&@\\x13'\n",
    "# Decrypted message: b'WTF, Alice!'\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "cb9d80de",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.7.7"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
